import numpy as np
import os

class AbaqusInpParser:
    def __init__(self, file_path,park_name,FRP_edge):
        self.file_path = file_path
        self.park_name = park_name
        self.FRP_edge = FRP_edge
        self._nodes = []           # 存储原始节点数据
        self._elements = []        # 存储原始单元数据
        self.current_set_labels = [] # 当前集合的节点标签

        # 最终输出的NumPy数组
        self.nodes_array = None    # 节点数组，格式：[[id, x, y], ...]
        self.elements_array = None # 单元数组
        self.origin_nodes_set_array = None  # 节点集合数组，格式：[[id, x, y], ...]
        self.nodes_set_array = None # 处理后的集合节点数组，格式：[[60000+id, x, y], ...]

        # 状态标志
        self._reading_nodes = False
        self._reading_elements = False
        self._reading_sets = False


    def get_set_nodes(self,offset=60000):
            """根据集合名称获取节点坐标数组"""
            # 获取节点ID列表
            node_ids = np.array(self.current_set_labels, dtype=np.int64)
            
            # 从nodes_array中提取（假设nodes_array格式为[[id, x, y], ...]）
            if self.nodes_array.size == 0:
                return np.array([])
                
            mask = np.isin(self.nodes_array[:, 0], node_ids)
            origin_nodes = self.nodes_array[mask].copy()  # 创建副本避免修改原始数据

            # 按X坐标排序（若X相同则按Y排序）
            sorted_indices = np.argsort(origin_nodes[:, 1])  # 仅按X排序
            sorted_origin = origin_nodes[sorted_indices].copy()

            # 对ID添加偏移量
            sorted_nodes = sorted_origin.copy()
            if offset != 0:
                sorted_nodes[:, 0] += offset  # 仅修改ID列

            # 保存排序结果
            self.origin_nodes_set_array = sorted_origin
            self.nodes_set_array = sorted_nodes

    def parse(self):
        """主方法，解析INP文件并提取节点和单元信息"""
        with open(self.file_path, 'r') as f:
            lines = f.readlines()
            current_part = False
            self._reset_reading_flags()
            #print(self._is_part_start(line=lines[2406].strip()))

            for line in lines:
                line = line.strip()
                if not line:
                    continue

                # 判断当前行是否为目标Part的开始或结束
                if self._is_part_start(line):
                    current_part = True
                    self._reset_reading_flags()
                    continue
                elif self._is_part_end(line) and current_part:
                    break

                # 如果当前行在目标Part内
                if current_part:
                    if self._is_node_section(line):
                        self._reset_reading_flags(keep_nodes=True)
                        continue
                    elif self._is_element_section(line):
                        self._reset_reading_flags(keep_elements=True)
                        continue
                    elif self._is_set_section(line):
                        self._reset_reading_flags()
                        if self._is_FRP_set(line):
                            self._reset_reading_flags(keep_sets=True)
                            continue
                        continue

                    # 提取节点、单元和集合信息
                    if self._reading_nodes:
                        self._parse_node_line(line)
                    elif self._reading_elements:
                        self._parse_element_line(line)
                    elif self._reading_sets:
                        self._parse_set_line(line)

        # 转换为NumPy数组
        nodes_array = np.array(self._nodes, dtype=np.float64) if self._nodes else np.array([])
        elements_array = np.array(self._elements, dtype=np.int64) if self._elements else np.array([])
        self.nodes_array = nodes_array
        self.elements_array = elements_array

#------------------------------------ 辅助方法 ------------------------------------------------#
    def _reset_reading_flags(self, keep_nodes=False, keep_elements=False, keep_sets=False):
        """重置读取状态标志"""
        self._reading_nodes = keep_nodes
        self._reading_elements = keep_elements
        self._reading_sets = keep_sets


 #---------------------------------- 解析判断方法 --------------------------------------------#
    def _is_part_start(self, line):
        """判断是否为Part的开始行"""
        return line.startswith(f'*Part, name={self.park_name}')

    def _is_part_end(self, line):
        """判断是否为Part的结束行"""
        return line.startswith('*End Part')

    def _is_node_section(self, line):
        """判断是否为节点部分的开始行"""
        return line.startswith('*Node')

    def _is_element_section(self, line):
        """判断是否为单元部分的开始行"""
        return line.startswith('*Element')
    
    def _is_set_section(self, line):
        """判断是否为集合部分的开始行"""
        return line.startswith('*Nset') or line.startswith('*Elset')
    
    def _is_FRP_set(self, line):
        """判断是否为FRP集合部分的开始行"""
        return line.startswith(f'*Nset, nset={self.FRP_edge}')

#------------------------------------ 数据解析方法 -------------------------------------#
    def _parse_node_line(self, line):
        """解析节点行"""
        parts = [p.strip() for p in line.split(',')]
        if len(parts) >= 3:
            try:
                node_id = int(parts[0])
                x = float(parts[1])
                y = float(parts[2])
                self._nodes.append([node_id, x, y])
            except ValueError:
                pass  # 忽略格式错误的行

    def _parse_element_line(self, line):
        """解析单元行"""
        parts = [p.strip() for p in line.split(',')]
        if len(parts) >= 5:
            try:
                elem_id = int(parts[0])
                elem_nodes = list(map(int, parts[1:5]))
                self._elements.append([elem_id] + elem_nodes)
            except ValueError:
                pass  # 忽略格式错误的行

    def _parse_set_line(self, line):
        """解析集合行,并得到行的坐标"""
        labels = list(map(int, line.replace(',', ' ').split()))
        self.current_set_labels.extend(labels)
    

#------------------------------------ 输出inp_with_cohesive -------------------------------------#

    def _format_nodes_set_array(self):
            """将nodes_set_array转换为INP格式字符串列表"""
            formatted = []
            for node in self.nodes_set_array:
                # 格式化ID为整数，坐标保留原有精度
                node_id = int(node[0])
                x = node[1]
                y = node[2]
                # 根据数值大小调整格式
                if abs(x) < 1e-10:   # 接近零时简化为"0."
                    x_str = "0."
                else:
                    # 保留最多8位小数，并删除末尾多余的零和小数点
                    x_str = f"{x:.8f}".rstrip('0').rstrip('.') if '.' in f"{x:.8f}" else f"{x:.0f}."
                    
                if abs(y) < 1e-10:
                    y_str = "0."
                else:
                    y_str = f"{y:.8f}".rstrip('0').rstrip('.') if '.' in f"{y:.8f}" else f"{y:.0f}."
                
                formatted.append(f"{node_id:>6d}, {x_str:>12s}, {y_str:>12s}\n")
            return formatted
    
    def _format_T2D2(self, start_elem_id=500001):
        """生成T2D2单元格式字符串列表（每行形如'101, 60001, 60002'）"""
        formatted = []
        elem_ids_T2D2 = []
        if self.nodes_set_array is None or len(self.nodes_set_array) < 2:
            return formatted
        
        # 提取偏移后的节点ID（确保按原始顺序排列）
        node_ids = self.nodes_set_array[:, 0].astype(int)
        
        # 每两个连续节点生成一个单元
        formatted.append("*Element, type=T2D2\n")
        for i, elem_id in enumerate(range(start_elem_id, start_elem_id + len(node_ids) - 1)):
            n1 = node_ids[i]
            n2 = node_ids[i+1]
            formatted.append(f"{elem_id:>6d}, {n1:>6d}, {n2:>6d}\n")
            elem_ids_T2D2.append(elem_id)  # 收集单元ID
        
        # 保存到实例变量以便后续访问
        self.elem_ids_T2D2 = elem_ids_T2D2
        return formatted
    
    def _format_T2D2_section(self, material_name="FRP", area=50.265, elset_name="FRP_BAR"):
        """生成T2D2单元的截面属性定义"""
        section = []
        if not hasattr(self, 'elem_ids_T2D2') or not self.elem_ids_T2D2:
            return section
        
        # 生成单元集合定义
        section.append(f"*Elset, elset={elset_name}, generate\n")
        section.append(f"{min(self.elem_ids_T2D2)}, {max(self.elem_ids_T2D2)}, 1\n")
        
        # 生成截面属性
        section.append(f"*Solid Section, elset={elset_name}, material={material_name}\n")
        section.append(f"{area:.2f},\n")  # 截面积参数
        return section

    def _format_COH2D4_section(self, material_name="Cohesive", thickness=50.265,
                               response_type="TRACTION SEPARATION", elset_name="FRP_COH2D4"):
        """生成COH2D4单元的粘聚层属性定义"""
        section = []
        if not hasattr(self, 'elem_ids_COH2D4') or not self.elem_ids_COH2D4:
            return section
        
        # 生成单元集合定义
        section.append(f"*Elset, elset={elset_name}, generate\n")
        section.append(f"{min(self.elem_ids_COH2D4)}, {max(self.elem_ids_COH2D4)}, 1\n")
        
        # 生成粘聚层属性
        section.append(
            f"*Cohesive Section, elset={elset_name}, "
            f"material={material_name}, response={response_type}, thickness=SPECIFIED\n"
        )
        section.append(f"1.,{thickness:.4f},\n")  # 粘聚层厚度
        return section

    def _format_COH2D4(self, start_elem_id=990001):
        """生成COH2D4单元格式字符串列表（每行形如'20001, 60001, 60002, 2, 1'）"""
        formatted = []
        elem_ids_COH2D4 = []
        if (self.origin_nodes_set_array is None or 
            self.nodes_set_array is None or 
            len(self.origin_nodes_set_array) != len(self.nodes_set_array)):
            return formatted
        
        # 提取原始和偏移后的节点ID
        origin_ids = self.origin_nodes_set_array[:, 0].astype(int)
        offset_ids = self.nodes_set_array[:, 0].astype(int)
        
        # 每对相邻节点生成一个单元
        formatted.append("*Element, type=COH2D4\n")
        for i, elem_id in enumerate(range(start_elem_id, start_elem_id + len(origin_ids) - 1)):
            off1 = offset_ids[i]
            off2 = offset_ids[i+1]
            orig1 = origin_ids[i+1]  # 注意原始节点顺序反向
            orig2 = origin_ids[i]
            formatted.append(f"{elem_id:>6d}, {off1:>6d}, {off2:>6d},{orig1:>6d}, {orig2:>6d}\n")
            elem_ids_COH2D4.append(elem_id)  # 收集单元ID
        
        # 保存到实例变量以便后续访问
        self.elem_ids_COH2D4 = elem_ids_COH2D4
        return formatted

    def insert_data(self, original_lines, insert_marker, prepend_data, offset_lines=0):
        """
        增强版数据插入方法，支持多标记判断
        :param insert_marker: 主标记或标记列表
        :param insert_marker1: 附加标记（可选）
        """
        # 构建标记列表
        markers = []
        if isinstance(insert_marker, (list, tuple)):
            markers.extend(insert_marker)
        else:
            markers.append(insert_marker)
    
        current_part = False
        # 查找所有匹配行位置
        marker_index  = None
        for i, line in enumerate(original_lines):
            # 判断当前行是否为目标Part的开始或结束
            if self._is_part_start(line):
                current_part = True
                continue
            elif self._is_part_end(line) and current_part:
                break

            if current_part:
                stripped = line.strip()
                for marker in markers:
                    if stripped.startswith(marker):
                        marker_index  = i
                        break  # 避免重复匹配
            if marker_index is not None:
                break  # 完全退出循环
    
        if marker_index is None:
            raise ValueError(f"未找到任何标记行: {markers}")
    
        modified_lines = original_lines.copy()
        
        insert_pos = marker_index  + offset_lines
        # 插入数据块（带前后空行）
        modified_lines[insert_pos:insert_pos] = prepend_data
    
        return modified_lines

    def output_inp(self):
        """将修改后的内容输出到同一目录下的新文件"""
        # 获取原始文件路径信息
        dir_path = os.path.dirname(self.file_path)     # 原始文件所在目录
        base_name = os.path.basename(self.file_path)   # 原始文件名
        new_filename = f"modified_{base_name}"         # 新文件名
        output_path = os.path.join(dir_path, new_filename)  # 完整输出路径

        # 读取原始内容
        with open(self.file_path, 'r') as f:
            original_lines = f.readlines()

        # 生成插入内容
        FRP_node = self._format_nodes_set_array()

        elemT2D2 = self._format_T2D2()
        elemCOH2D4 = self._format_COH2D4()

        T2D2_section = self._format_T2D2_section()
        COH2D4_section = self._format_COH2D4_section()

        # 插入数据
        modified_lines = self.insert_data(
            original_lines=original_lines,
            insert_marker='*Element',
            prepend_data=FRP_node
        )
        modified_lines = self.insert_data(
            original_lines=modified_lines,
            insert_marker=['*Nset','*Elset'],
            prepend_data=elemT2D2
        )
        modified_lines = self.insert_data(
            original_lines=modified_lines,
            insert_marker=['*Nset','*Elset'],
            prepend_data=elemCOH2D4
        )
        modified_lines = self.insert_data(
            original_lines=modified_lines,
            insert_marker=['*Nset','*Elset'],
            prepend_data=COH2D4_section
        )
        modified_lines = self.insert_data(
            original_lines=modified_lines,
            insert_marker=['*Nset','*Elset'],
            prepend_data=T2D2_section
        )
        # 写入新文件（保留原始换行符）
        with open(output_path, 'w', newline='') as f:
            f.writelines(modified_lines)


    def main(self):
        self.parse()
        self.get_set_nodes()
        self.output_inp()
            
if __name__ == "__main__":
    inp_addre = r"F:\000000毕业论文撰写\6第六章\inp\tunnel_cohesive.inp"
    parser = AbaqusInpParser(file_path=inp_addre,park_name='tunnel',FRP_edge='FRP_edge')
    parser.main()
    #print(parser.origin_nodes_set_array)
        

